import math
import struct
import inspect
import time

from Config import LLMConfig
from typing import Any,Optional,Tuple,List
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.eps = eps #防止除0设置的小数
        self.weight = nn.Parameter(torch.ones(dim)) #可学习的参数矩阵

    def forward(self, x):# x除以 dim维度的均方根 再乘以可学习的参数矩阵 最后调整为x的数据类型返回
        return self.weight * (x.float() * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)

#预设置的旋转位置编码里面三角函数内的参数，固定不变
def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
    #dim维度中 2i分量对应参数1e6^(-2i/d)  2i+1分量对应参数1e6^((-2i+1)/d)
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) 
    t = torch.arange(end, device=freqs.device)  #位置下标
    freqs = torch.outer(t, freqs).float()  #把位置下标和dim对应参数相乘 得到最终三角函数内参数
    pos_cis = torch.polar(torch.ones_like(freqs), freqs)  #转换为极坐标,引入三角函数
    return pos_cis 
    
def apply_rotary_emb(xq, xk, pos_cis):
    def unite_shape(pos_cis, x): #
        ndim = x.ndim
        assert 0 <= 1 < ndim
        assert pos_cis.shape == (x.shape[1], x.shape[-1]) #pos_cis形状应该是(seq_len, head_dim)
        shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] #规定一个shape：x的第一个维度(seq_len)和最后一个维度(head_dim)保持不变，其余设置为1
        return pos_cis.view(*shape) #把规定的shape应用于pos_cis
    
    #使用 reshape 方法将最后一个维度拆分为两个维度，拆完后的倒数第二个维度减半，倒数第一个维度值为2
    #再使用 torch.view_as_complex 方法将最后两个拆分的维度合并为一个复数维度，但最后一维数量减半。 所以 xq_ 形状是(bsz , seq_len , n_heads , head_dim//2)
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) 
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    pos_cis = unite_shape(pos_cis, xq_)
    #位置编码和q k相乘并展成实数
    xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk) 
    
    
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    bs, slen, n_kv_heads, head_dim = x.shape
    if n_rep == 1:
        return x

    return (
        x[:, :, :, None, :]    #在原先第2个(从0开始)维度后面插入一个维度 用于重复n_rep次 
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim) #再展成原先x形状 
    )

class Attention(nn.Module):
    def __init__(self, args: LLMConfig):
        super().__init__()
        #这里分成n_heads 与 n_kv_heads不一致 是因为用了GQA n_heads对应query数量 n_kv_heads对应key value数量
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads #设置了n_kv_heads就使用GQA 否则不使用
        assert args.n_heads % self.n_kv_heads == 0
        self.n_heads = args.n_heads 
        self.n_rep = self.n_heads // self.n_kv_heads
        self.head_dim = args.dim // args.n_heads
        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
        self.attn_dropout = nn.Dropout(args.dropout)
        self.resid_dropout = nn.Dropout(args.dropout)
        mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf")) #注意力掩码矩阵(max_seq_len,max_seq_len)
        mask = torch.triu(mask, diagonal=1) #下三角包括主对角线为0 其余负无穷
        self.register_buffer("mask", mask, persistent=False) 
    

    def forward(self,
                x: torch.Tensor,
                pos_cis: torch.Tensor,
                past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
                use_cache=False):
        bsz, seq_len, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
        xq = xq.view(bsz, seq_len, self.n_heads, self.head_dim)
        xk = xk.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
        xv = xv.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
        xq, xk = apply_rotary_emb(xq, xk, pos_cis)
        # kv_cache实现
        if past_key_value is not None: #只有推理模式才会启用
            xk = torch.cat([past_key_value[0], xk], dim=1)
            xv = torch.cat([past_key_value[1], xv], dim=1)
            #这里 K V拼接后seq_len维度变为所有以前生成序列长度 和 Q的seq_len(=1)不一致
        past_kv = (xk, xv) if use_cache else None
        xq, xk, xv = ( #形状变为 (bsz, n_heads, seq_len, head_dim)
            xq.transpose(1, 2), 
            repeat_kv(xk, self.n_rep).transpose(1, 2),
            repeat_kv(xv, self.n_rep).transpose(1, 2)
        )
        scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) #四维矩阵做矩阵乘，只对最后两维进行矩阵运算，前两维度用来批量计算
        scores += self.mask[:, :, :seq_len, :seq_len] #这里的seq_len在推理模式=1 训练模式=文本最大长度
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        scores = self.attn_dropout(scores)
        output = scores @ xv 

        output = output.transpose(1, 2).reshape(bsz, seq_len, -1) #形状变为 (bsz, seq_len, n_heads*head_dim) 也就是(bsz, seq_len, dim)
        output = self.resid_dropout(self.wo(output))
        return output, past_kv

        

class FeedForward(nn.Module):
    def __init__(self, config: LLMConfig):
        super().__init__()
        if config.hidden_dim is None:
            hidden_dim = 4 * config.dim
            hidden_dim = int(2 * hidden_dim / 3)
            config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
        self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)
        self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)
        self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)
        self.dropout = nn.Dropout(config.dropout)
     
    def forward(self, x):
        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
     
class MyGOBlock(nn.Module):
    def __init__(self, layer_id:int ,config: LLMConfig):
          super().__init__()
          self.n_heads = config.n_heads
          self.dim = config.dim
          self.head_dim = config.dim // config.n_heads
          self.attention = Attention(config)
          self.layer_id=layer_id
          self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
          self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
          self.feed_forward = FeedForward(config)
     
    def forward(self, x, pos_cis, past_key_value=None, use_cache=False):
          h_attn, past_kv = self.attention(
            self.attention_norm(x),
            pos_cis,
            past_key_value=past_key_value,
            use_cache=use_cache
          )
          h = x + h_attn 
          out = h + self.feed_forward(self.ffn_norm(h))
          return out, past_kv

class MyGO(PreTrainedModel):
    config_class = LLMConfig

    def __init__(self, params:LLMConfig=None):
        self.params = params or LLMConfig()
        super().__init__(params)
        self.vocab_size , self.n_layers= params.vocab_size, params.n_layers
        self.tok_embeddings = nn.Embedding(self.vocab_size, params.dim)
        self.dropout = nn.Dropout(params.dropout)
        self.layers = nn.ModuleList([MyGOBlock(l, params) for l in range(self.n_layers)])
        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
        self.tok_embeddings.weight = self.output.weight
        self.register_buffer("pos_cis",
                             precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),
                             persistent=False)
        self.OUT = CausalLMOutputWithPast()

    def forward(self,
                 input_ids: Optional[torch.Tensor] = None,
                 past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
                 use_cache: bool = False,
                 **args
                 ):
        # KVcache分层存储，若为空则置空len(self.layers)层
        past_key_values = past_key_values or [None] * len(self.layers)
        #开始位置，没设置则默认0
        start_pos = args.get('start_pos', 0)
        #把输入token_id通过embedding映射到dim维度
        h = self.dropout(self.tok_embeddings(input_ids))
        pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)] #获取位置编码中三角函数内参数
        past_kvs=[]

        for l, layer in enumerate(self.layers):
            h, past_kv = layer(
                h, pos_cis,
                past_key_value=past_key_values[l],
                use_cache=use_cache
            )
            past_kvs.append(past_kv) 
        logits=self.output(self.norm(h)) #logits尺寸是 (bsz, seq_len, vocab_size) 
        self.OUT.__setitem__('logits', logits)
        self.OUT.__setitem__('past_key_values', past_kvs)

        return self.OUT
    
    #推理模式
    @torch.inference_mode()
    def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,
                stream=False, rp=1., use_cache=True, pad_token_id=0, **args):
            
        return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)

           
    #这里input_ids应该是一个一维序列token_id组成向量，但模型中输入input_ids要求有batch_size维度 所以在最前面加上一个1维度，变成(1, len(input_ids))
    def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args):
            start, first_seq, past_kvs = input_ids.shape[1], True, None
            while input_ids.shape[1] < max_new_tokens - 1:
                if first_seq or not use_cache:
                    out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False
                else:
                    out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache,
                        start_pos=input_ids.shape[1] - 1, **args)
                logits, past_kvs = out.logits[:, -1, :], out.past_key_values  #只取seq_len维度中最后一个元素对应各个vocab_size中概率分数
                logits[:, list(set(input_ids.tolist()[0]))] /= rp #降低已经生成token分数，降低再生成的概率
                logits /= (temperature + 1e-9) #应用温度值，使分布更加平滑
                #核采样/top-p 取累积最高概率到p的token_id进一步计算softmax，根据概率取最终token
                if top_p is not None and top_p < 1.0: 
                    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
                    sorted_probs = F.softmax(sorted_logits, dim=-1)
                    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
                    sorted_indices_to_remove = cumulative_probs > top_p
                    sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone() #把掩码整体右移一位，防止掩盖最大概率的token
                    sorted_indices_to_remove[:, 0] = False #保留一个最大概率的token
                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) 
                    logits[indices_to_remove] = -float('Inf')
                #根据采样后的词计算概率并根据概率抽样，作为最终新的token
                input_ids_next = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
                input_ids = torch.cat((input_ids, input_ids_next), dim=1)
                #返回一个生成器
                yield input_ids[:, start:]
                #这里尺寸必须是1*1，否则不是标量，会报错
                if input_ids_next.item() == eos_token_id:
                    break
